import os
import pandas as pd
from openai import OpenAI
from langchain.prompts import PromptTemplate
from langchain.schema import BaseOutputParser
from langchain.chains import LLMChain
from langchain_openai import ChatOpenAI
from typing import List, Dict
import json
import re
import glob
import shutil
from dotenv import load_dotenv

load_dotenv()

# 阿里云百炼配置
DASHSCOPE_CONFIG = {
    "api_key": os.getenv('DASHSCOPE_API_KEY'),
    "base_url": os.getenv('BASE_URL'),
    "model": "qwen-long"
}

# 验证API密钥
if not DASHSCOPE_CONFIG["api_key"]:
    raise ValueError("请设置DASHSCOPE_API_KEY环境变量")

class QAOutputParser(BaseOutputParser):
    """解析LLM输出的问答对"""

    def parse(self, text: str) -> List[Dict[str, str]]:
        """解析问答对文本"""
        qa_pairs = []

        # 预处理：提取JSON内容（处理代码块格式）
        json_text = self._extract_json_content(text)

        # 尝试解析JSON格式
        try:
            if json_text.strip().startswith('[') or json_text.strip().startswith('{'):
                data = json.loads(json_text)
                if isinstance(data, list):
                    return data
                elif isinstance(data, dict) and 'qa_pairs' in data:
                    return data['qa_pairs']
        except json.JSONDecodeError:
            pass

        # 解析文本格式的问答对
        lines = text.strip().split('\n')
        current_qa = {}

        for line in lines:
            line = line.strip()
            if not line:
                continue

            # 匹配问题格式 - 提取纯净的问题内容
            q_match = re.match(r'^(?:问题?\d*|Q\d*|Question\d*)\s*[：:：]?\s*(.+)$', line, re.IGNORECASE)
            if q_match:
                if current_qa and 'question' in current_qa and 'answer' in current_qa:
                    qa_pairs.append(current_qa)
                # 清理问题内容，移除可能的标号和冒号
                clean_question = self._clean_content(q_match.group(1))
                current_qa = {'question': clean_question}
                continue

            # 匹配答案格式 - 提取纯净的答案内容
            a_match = re.match(r'^(?:答案?\d*|A\d*|Answer\d*)\s*[：:：]?\s*(.+)$', line, re.IGNORECASE)
            if a_match and 'question' in current_qa:
                # 清理答案内容，移除可能的标号和冒号
                clean_answer = self._clean_content(a_match.group(1))
                current_qa['answer'] = clean_answer
                continue

            # 如果当前有问题但没有答案，将此行作为答案
            if 'question' in current_qa and 'answer' not in current_qa:
                clean_answer = self._clean_content(line)
                current_qa['answer'] = clean_answer

        # 添加最后一个问答对
        if current_qa and 'question' in current_qa and 'answer' in current_qa:
            qa_pairs.append(current_qa)

        return qa_pairs

    def _extract_json_content(self, text: str) -> str:
        """提取JSON内容，处理代码块格式"""
        text = text.strip()

        # 检查是否包含```json代码块
        json_block_pattern = r'```json\s*\n?([\s\S]*?)\n?```'
        json_match = re.search(json_block_pattern, text, re.IGNORECASE)

        if json_match:
            # 提取代码块中的JSON内容
            return json_match.group(1).strip()

        # 如果没有代码块，返回原文本
        return text

    def _clean_content(self, content: str) -> str:
        """清理内容，移除标号和多余的冒号"""
        # 移除开头的数字标号和冒号
        content = re.sub(r'^\d+[.、：:：]?\s*', '', content)
        # 移除开头的字母标号和冒号
        content = re.sub(r'^[A-Za-z][.、：:：]?\s*', '', content)
        # 移除多余的冒号
        content = re.sub(r'^[：:：]+\s*', '', content)
        return content.strip()


class CSVQAGenerator:
    """CSV问答对生成器"""

    def __init__(self):
        # 初始化阿里云百炼模型
        self.llm = ChatOpenAI(
            api_key=DASHSCOPE_CONFIG["api_key"],
            base_url=DASHSCOPE_CONFIG["base_url"],
            model=DASHSCOPE_CONFIG["model"],
            temperature=0.7
        )

        # 创建标题行识别的提示模板
        self.header_detection_template = PromptTemplate(
            input_variables=["table_rows"],
            template="""请帮我分析下面这个表格的结构。标题行是指用于描述下方数据列含义的行，通常不是数据本身。

表格数据：
{table_rows}

请遵循以下步骤：
1. 仔细观察表格的结构。
2. 判断从第1行开始，连续有几行是作为整个表格的标题行（表头）。
3. 回答格式为：标题行数量：X行
"""
        )

        # 创建标题行识别LLM链
        self.header_detection_chain = LLMChain(
            llm=self.llm,
            prompt=self.header_detection_template
        )

        # 创建表格问答生成的提示模板（包含表名）
        self.table_prompt_template = PromptTemplate(
            input_variables=["table_name", "header", "row_data"],
            template="""你是一个专业的问答生成专家。请根据以下表格数据，生成中文问答对。

表名：
{table_name}

表格标题行：
{header}

数据行：
{row_data}

要求：
1. 基于表格标题和数据行的对应关系生成问答对
2. 根据数据内容的丰富程度，决定问答对生成数量，涵盖该行数据的关键信息
3. 问答对必须基于提供的数据，不能添加或修改任何信息
4. 主语不能使用"这""那"等代词,问题中不要出现表名
5. 行列对应限制信息都要说明

请以标准JSON格式输出问答对：

{{
  "qa_pairs": [
    {{
      "question": " ",
      "answer": " "
    }},
    {{
      "question": " ",
      "answer": " "
    }}
  ]
}}

注意：
1. 只输出JSON格式的结果
2. 确保JSON格式正确，使用双引号
3. 问题和答案内容中的引号需要转义
4. 不要输出其他格式的内容
"""
        )

        # 创建表格问答生成LLM链
        self.table_qa_chain = LLMChain(
            llm=self.llm,
            prompt=self.table_prompt_template,
            output_parser=QAOutputParser()
        )

    def detect_header_rows(self, table_rows: List[str]) -> int:
        """使用大模型检测标题行数量"""
        try:
            # 取前5行进行分析
            analysis_rows = table_rows[:5]
            table_text = "\n".join([f"第{i+1}行: {row}" for i, row in enumerate(analysis_rows)])
            
            # 调用大模型分析
            result = self.header_detection_chain.run(table_rows=table_text)
            # print(result)
            
            # 优先从"标题行数量：X行"格式中提取数字
            pattern_match = re.search(r'标题行数量[：:：]?\s*(\d+)\s*行?', result)
            if pattern_match:
                header_count = int(pattern_match.group(1))
            else:
                # 如果没有找到标准格式，尝试提取最后出现的数字
                numbers = re.findall(r'\d+', result)
                if numbers:
                    header_count = int(numbers[-1])  # 取最后一个数字
                else:
                    header_count = 1
            
            # 确保结果在合理范围内
            header_count = max(1, min(header_count, len(table_rows) - 1))
            
            print(f"    检测到标题行数量: {header_count}")
            return header_count
            
        except Exception as e:
            print(f"标题行检测失败，使用默认值1: {e}")
            return 1

    def detect_table_name(self, first_row: str) -> tuple:
        """检测表格第一行是否为表名，并判断表格类型"""
        try:
            # 分割第一行的单元格
            cells = [cell.strip() for cell in first_row.split(',')]
            
            # 检查第二个单元格是否为空
            if len(cells) >= 2:
                second_cell_empty = not cells[1]  # 第二个单元格是否为空
                
                if second_cell_empty:
                    # 第二个单元格为空，这是类型一表格（需要处理）
                    if cells[0]:
                        # 第一个单元格有内容，认为是表名
                        return cells[0], True  # (表名, 是否处理)
                    else:
                        # 第一个单元格也为空，没有表名但仍需处理
                        return "", True
                else:
                    # 第二个单元格不为空，这是类型二表格（跳过处理）
                    return "", False
            elif len(cells) == 1 and cells[0]:
                # 只有一个单元格，且有内容，认为是表名，需要处理
                return cells[0], True
            else:
                # 其他情况，按需要处理的表格处理
                return "", True
                
        except Exception as e:
            print(f"检测表名时出错: {e}")
            return "", True  # 出错时默认处理

    def generate_qa_from_table_row(self, header: str, row_data: str, header_rows_count: int = 1, table_name: str = "") -> List[Dict[str, str]]:
        """根据表格标题行和数据行生成问答对"""
        try:
            # 如果没有标题行（header_rows_count为0），则只使用数据行生成问答对
            if header_rows_count == 0:
                # 创建无标题行的提示模板（包含表名）
                no_header_template = PromptTemplate(
                    input_variables=["table_name", "row_data"],
                    template="""你是一个专业的问答生成专家。请根据以下数据行，生成中文问答对。

表名：
{table_name}

数据行：
{row_data}

要求：
1. 基于数据行的内容生成问答对
2. 根据数据内容的丰富程度，决定问答对生成数量，涵盖该行数据的关键信息
3. 问答对必须基于提供的数据，不能添加或修改任何信息
4. 主语不能使用"这""那"等代词，问题中不要出现表名。
5. 行列对应限制信息都要说明。

请以标准JSON格式输出问答对：

{{
  "qa_pairs": [
    {{
      "question": " ",
      "answer": " "
    }},
    {{
      "question": " ",
      "answer": " "
    }}
  ]
}}

注意：
1. 只输出JSON格式的结果
2. 确保JSON格式正确，使用双引号
3. 问题和答案内容中的引号需要转义
4. 不要输出其他格式的内容
"""
                )
                
                # 创建无标题行的LLM链
                no_header_chain = LLMChain(
                    llm=self.llm,
                    prompt=no_header_template,
                    output_parser=QAOutputParser()
                )
                
                result = no_header_chain.run(table_name=table_name or "未知表格", row_data=row_data)
            else:
                # 有标题行时使用原有逻辑
                result = self.table_qa_chain.run(table_name=table_name or "未知表格", header=header, row_data=row_data)
            
            return result if isinstance(result, list) else []
        except Exception as e:
            print(f"生成表格问答对时出错: {e}")
            return []

    def parse_csv_tables(self, csv_file_path: str) -> List[Dict]:
        """解析CSV文件中以空行分隔的表格"""
        tables = []
        
        try:
            # 尝试多种编码格式读取CSV文件
            encodings = ['utf-8', 'gbk', 'gb2312', 'utf-8-sig', 'latin1']
            lines = None
            
            for encoding in encodings:
                try:
                    with open(csv_file_path, 'r', encoding=encoding) as f:
                        lines = f.readlines()
                    print(f"成功使用 {encoding} 编码读取文件")
                    break
                except UnicodeDecodeError:
                    continue
            
            if lines is None:
                print(f"无法使用任何编码格式读取文件: {csv_file_path}")
                return tables
            
            current_table = []
            
            for line in lines:
                line = line.strip()
                
                # 检查是否为空行或只包含分隔符的行
                is_empty_line = not line or re.match(r'^[,\s]*$', line)
                
                if is_empty_line:
                    if current_table:
                        # 处理当前表格
                        if len(current_table) >= 1:
                            # 检测表名和表格类型
                            table_name, should_process = self.detect_table_name(current_table[0])
                            
                            if not should_process:
                                # 跳过类型二表格（第一行第二个单元格不为空）
                                print(f"    跳过类型二表格（第一行第二个单元格不为空）")
                                current_table = []
                                continue
                            
                            print(f"    检测到表名: {table_name}")
                            
                            # 如果有表名，从第二行开始检测标题行（完全跳过第一行）
                            if table_name:
                                # 只传输第2-6行给大模型进行标题行检测，完全不包含第一行
                                header_detection_rows = current_table[1:6]  # 第2-6行
                                if header_detection_rows:
                                    header_rows_count = self.detect_header_rows(header_detection_rows)
                                    # 标题行从第2行开始计算
                                    header_rows = current_table[1:1+header_rows_count]  # 第2行开始的标题行
                                    data_rows = current_table[1+header_rows_count:]     # 数据行
                                else:
                                    header_rows_count = 0
                                    header_rows = []
                                    data_rows = current_table[1:]  # 跳过表名行
                            else:
                                # 没有表名，使用原有逻辑
                                header_rows_count = self.detect_header_rows(current_table)
                                header_rows = current_table[:header_rows_count]
                                data_rows = current_table[header_rows_count:]
                            
                            # 合并多行标题为一个标题字符串
                            header = "\n".join(header_rows)
                            
                            if data_rows:  # 确保有数据行
                                tables.append({
                                    'table_name': table_name,
                                    'header': header,
                                    'data_rows': data_rows,
                                    'header_rows_count': header_rows_count
                                })
                        current_table = []
                else:
                    # 过滤掉row_count和col_count信息及其后面的数字
                    # 使用正则表达式匹配并删除这些信息
                    filtered_line = re.sub(r'\brow_count:\s*\d+', '', line)
                    filtered_line = re.sub(r'\bcol_count:\s*\d+', '', filtered_line)
                    # 清理多余的空格和逗号
                    filtered_line = re.sub(r'\s*,\s*,', ',', filtered_line)  # 去除连续逗号
                    filtered_line = re.sub(r'^\s*,\s*', '', filtered_line)    # 去除开头的逗号
                    filtered_line = re.sub(r'\s*,\s*$', '', filtered_line)    # 去除结尾的逗号
                    filtered_line = filtered_line.strip()
                    
                    if filtered_line:  # 只有非空行才添加
                        current_table.append(filtered_line)
            
            # 处理最后一个表格（如果文件末尾没有空行）
            if current_table and len(current_table) >= 1:
                # 检测表名和表格类型
                table_name, should_process = self.detect_table_name(current_table[0])
                
                if not should_process:
                    # 跳过类型二表格（第一行第二个单元格不为空）
                    print(f"    跳过类型二表格（第一行第二个单元格不为空）")
                else:
                    print(f"    检测到表名: {table_name}")
                    
                    # 如果有表名，从第二行开始检测标题行（完全跳过第一行）
                    if table_name:
                        # 只传输第2-6行给大模型进行标题行检测，完全不包含第一行
                        header_detection_rows = current_table[1:6]  # 第2-6行
                        if header_detection_rows:
                            header_rows_count = self.detect_header_rows(header_detection_rows)
                            # 标题行从第2行开始计算
                            header_rows = current_table[1:1+header_rows_count]  # 第2行开始的标题行
                            data_rows = current_table[1+header_rows_count:]     # 数据行
                        else:
                            header_rows_count = 0
                            header_rows = []
                            data_rows = current_table[1:]  # 跳过表名行
                    else:
                        # 没有表名，使用原有逻辑
                        header_rows_count = self.detect_header_rows(current_table)
                        header_rows = current_table[:header_rows_count]
                        data_rows = current_table[header_rows_count:]
                    
                    # 合并多行标题为一个标题字符串
                    header = "\n".join(header_rows)
                    
                    if data_rows:  # 确保有数据行
                        tables.append({
                            'table_name': table_name,
                            'header': header,
                            'data_rows': data_rows,
                            'header_rows_count': header_rows_count
                        })
                
        except Exception as e:
            print(f"解析CSV文件时出错: {e}")
            
        return tables

    def process_csv_file(self, csv_file_path: str, output_file: str):
        """处理单个CSV文件"""
        print(f"正在处理文件: {os.path.basename(csv_file_path)}")
        
        # 解析CSV文件中的表格
        tables = self.parse_csv_tables(csv_file_path)
        
        if not tables:
            print(f"文件 {os.path.basename(csv_file_path)} 中未找到有效表格")
            return
        
        print(f"找到 {len(tables)} 个表格")
        
        all_qa_pairs = []
        table_index = 0
        
        for table in tables:
            table_index += 1
            table_name = table.get('table_name', '')  # 从table字典中获取table_name
            header = table['header']
            data_rows = table['data_rows']
            header_rows_count = table.get('header_rows_count', 1)
            
            table_info = f"表格 {table_index}"
            if table_name:
                table_info += f" (表名: {table_name})"
            table_info += f": 标题行数={header_rows_count}, 数据行数={len(data_rows)}"
            print(f"  处理{table_info}")
            
            for row_index, row_data in enumerate(data_rows, 1):
                # print(f"    处理第 {row_index} 行数据...")
                
                # 生成问答对，传递表名参数
                qa_pairs = self.generate_qa_from_table_row(header, row_data, header_rows_count, table_name)
                
                if qa_pairs:
                    # 为每个问答对添加表格和行信息
                    for qa_index, qa in enumerate(qa_pairs, 1):
                        qa['table_index'] = table_index
                        qa['row_index'] = row_index
                        qa['qa_index_in_row'] = qa_index
                        # 添加表名信息
                        if table_name:
                            qa['table_name'] = table_name
                        # 只有当有标题行时才添加标题信息
                        if header_rows_count > 0:
                            qa['table_header'] = header
                        qa['row_data'] = row_data
                        qa['header_rows_count'] = header_rows_count
                    
                    all_qa_pairs.extend(qa_pairs)
                    # print(f"      成功生成 {len(qa_pairs)} 个问答对")
                else:
                    print(f"      第 {row_index} 行未能生成问答对")
        
        # 保存到Excel文件
        if all_qa_pairs:
            category_id = os.path.splitext(os.path.basename(csv_file_path))[0]
            self.save_to_excel(all_qa_pairs, output_file, category_id)
            print(f"文件 {os.path.basename(csv_file_path)} 处理完成，共生成 {len(all_qa_pairs)} 个问答对")
        else:
            print(f"文件 {os.path.basename(csv_file_path)} 未生成任何问答对")

    def process_csv_directory(self, input_dir: str, output_dir: str):
        """批量处理目录中的所有CSV文件"""
        # 查找所有csv文件
        csv_pattern = os.path.join(input_dir, "*.csv")
        csv_files = glob.glob(csv_pattern)
        
        if not csv_files:
            print(f"在目录 {input_dir} 中未找到任何.csv文件")
            return
        
        print(f"找到 {len(csv_files)} 个CSV文件")
        
        # 确保输出目录存在
        os.makedirs(output_dir, exist_ok=True)
        
        total_qa_pairs = 0
        
        for i, csv_file in enumerate(csv_files, 1):
            print(f"\n正在处理第 {i}/{len(csv_files)} 个文件: {os.path.basename(csv_file)}")
            
            # 生成输出文件名
            base_name = os.path.splitext(os.path.basename(csv_file))[0]
            output_file = os.path.join(output_dir, f"{base_name}a.xlsx")
            
            # 如果输出文件已存在，先删除
            if os.path.exists(output_file):
                os.remove(output_file)
            
            try:
                # 处理CSV文件
                self.process_csv_file(csv_file, output_file)
                
                # 统计生成的问答对数量
                if os.path.exists(output_file):
                    df = pd.read_excel(output_file)
                    file_qa_count = len(df)
                    total_qa_pairs += file_qa_count
                    print(f"文件 {base_name}a.xlsx 已保存，包含 {file_qa_count} 个问答对")
                
            except Exception as e:
                print(f"处理文件 {os.path.basename(csv_file)} 时出错: {e}")
                continue
        
        print(f"\n批量处理完成！")
        print(f"总共处理了 {len(csv_files)} 个CSV文件")
        print(f"总共生成了 {total_qa_pairs} 个问答对")
        print(f"结果已保存到目录: {output_dir}")

    def save_to_excel(self, qa_pairs: List[Dict[str, str]], filename: str, category_id: str = None):
        """保存问答对到Excel文件"""
        if not qa_pairs:
            print("没有问答对可保存")
            return
        
        # 准备数据
        data = []
        for i, qa in enumerate(qa_pairs, 1):
            row = {
                '标号': i,
                '提问': qa.get('question', ''),
                '回答': qa.get('answer', '')
            }
            if category_id:
                row['category_id'] = category_id
            # 只有当存在标题行时才添加表格标题信息
            if 'table_header' in qa and qa.get('header_rows_count', 1) > 0:
                row['表格标题'] = qa['table_header']
            if 'row_data' in qa:
                row['原始数据'] = qa['row_data']
            data.append(row)
        
        # 创建DataFrame，只包含基本列和必要信息（不包含序号相关列）
        columns = ['标号', '提问', '回答']
        if category_id:
            columns.append('category_id')
        if any('表格标题' in row for row in data):
            columns.append('表格标题')
        if any('原始数据' in row for row in data):
            columns.append('原始数据')
        
        df = pd.DataFrame(data, columns=columns)
        df.to_excel(filename, index=False)


def clear_folder(folder_path: str):
    """清空指定文件夹中的所有文件"""
    if os.path.exists(folder_path):
        try:
            shutil.rmtree(folder_path)
            os.makedirs(folder_path, exist_ok=True)
            print(f"已清空文件夹: {folder_path}")
        except Exception as e:
            print(f"清空文件夹失败 {folder_path}: {e}")
    else:
        os.makedirs(folder_path, exist_ok=True)
        print(f"已创建文件夹: {folder_path}")

def main():
    """主函数 - 批量处理CSV文件"""
    generator = CSVQAGenerator()
    
    print("=== CSV问答对生成工具 ===") 
    print(f"连接到阿里云百炼服务: {DASHSCOPE_CONFIG['base_url']}")
    print(f"使用模型: {DASHSCOPE_CONFIG['model']}")
    print()
    
    # 设置输入目录和输出目录
    current_dir = os.path.dirname(os.path.abspath(__file__))
    input_dir = os.path.join(os.path.dirname(current_dir), "A_segmenter", "csv")
    output_dir = os.path.join(current_dir, "csvQAa")
    
    print(f"当前工作目录: {current_dir}")
    print(f"输入目录: {input_dir}")
    print(f"输出目录: {output_dir}")
    print()
    
    # 检查输入目录是否存在
    if not os.path.exists(input_dir):
        print(f"错误：输入目录不存在 - {input_dir}")
        return
    
    # 清空输出目录
    clear_folder(output_dir)
    
    try:
        # 批量处理目录中的所有CSV文件
        generator.process_csv_directory(input_dir, output_dir)
        
    except Exception as e:
        print(f"批量处理过程中发生错误: {e}")


if __name__ == "__main__":
    main()